import io
from typing import Optional, Union
from chernc.treediff.gumtree.gen.tree_generator import TreeGenerator
from chernc.treediff.gumtree.tree.tree_context import TreeContext
from tree_sitter import Language, Parser, Tree

from chernc.treediff.gumtree.gen.generators.treesitter_tree import TreeSitterNode


class TreeSitterTreeGenerator(TreeGenerator):
    def __init__(self, language: str):
        """
        Initialize TreeSitterTreeGenerator with a specific language.

        :param language: The language used for parsing (e.g., "java", "python").
        """
        
        # 加载tree-sitter语言的库文件
        if language == "rust":
            import tree_sitter_rust as tsrust
            self.language = Language(tsrust.language())
            self.parser = Parser(self.language)
        else:
            raise Exception(f"Invalid language name {language}...")

    def generate(self, r: Union[io.TextIOWrapper, str]) -> TreeContext:
        """
        Generate a TreeContext from the input stream using TreeSitter.

        :param r: The input stream (e.g., a file reader).
        :return: A TreeContext generated from the input.
        """
        if isinstance(r, io.TextIOWrapper):
            content = r.read()  # 读取整个输入流的内容
        elif isinstance(r, str):
            content = r
        tree = self.parser.parse(content.encode(), encoding="utf8")  # 使用TreeSitter解析代码内容
        return self.convert_to_tree_context(tree)

    def convert_to_tree_context(self, tree: Tree) -> TreeContext:
        """
        Convert the TreeSitter parsed tree to a GumTree TreeContext.

        :param tree: The parsed tree from TreeSitter.
        :return: A TreeContext object representing the parsed tree.
        """
        context = TreeContext()
        context.set_root(TreeSitterNode(tree.root_node))
        return context
